
/**
 * @license
 * Copyright 2019 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
import {neg} from '@tensorflow/tfjs-core';
import {NamedTensorsMap} from '../data/types';

import {Graph, Node} from '../operations/types';

/**
 * This graph rewrite rule tries to identify the PRelu structure generated by
 * tf.keras, and convert it to tfjs core prelu op.
 *
 * The formula of PReLU is:
 * f(x) = alpha * x for x < 0, f(x) = x for x >= 0.
 *
 * `x` is the input, and `alpha` is a trainable tensor which can be broadcasted
 * to the shape of `x`.
 *
 * There's no native PRelu op in TensorFlow, so tf.keras generates the following
 * structure which does the equivalent calculation:
 * f(x) = Relu(x) + (-alpha * Relu(-x))
 *
 * Practically, alpha is always a constant in the inference graph.
 * Therefore, we're looking for the structure:
 *
 * f(x) = Add(Relu(x), Mul(negative_alpha, Relu(Neg(x))))
 *
 * And generate the follow sub graph:
 * f(x) = Prelu(x, neg(negative_alpha))
 *
 * @param graph Graph, model graph object
 * @param weightMap NamedTensorsMap, the weight map for the executor.
 */

export function rewritePrelu(graph: Graph, weightMap: NamedTensorsMap) {
  for (const key in graph.nodes) {
    const addNode = graph.nodes[key];
    if (addNode == null || addNode.op !== 'Add' && addNode.op !== 'AddV2' ||
        addNode.inputNames.length !== 2) {
      continue;
    }

    const reluNode = addNode.inputs.find(input => input.op === 'Relu');
    if (reluNode == null || reluNode.inputNames.length !== 1) {
      continue;
    }

    const mulOp = addNode.inputs.find(input => input.op === 'Mul');
    if (mulOp == null || mulOp.inputNames.length !== 2) {
      continue;
    }

    const negAlphaTensorNode = mulOp.inputs.find(input => input.op === 'Const');

    const reluNegInputNode = mulOp.inputs.find(input => input.op === 'Relu');

    if (negAlphaTensorNode == null || reluNegInputNode == null ||
        reluNegInputNode.inputNames.length !== 1) {
      continue;
    }

    // This detects a Neg op followed by a separated Relu op.
    const negInputNode = reluNegInputNode.inputs[0];
    if (negInputNode == null || negInputNode.op !== 'Neg' ||
        negInputNode.inputNames.length !== 1) {
      continue;
    }

    if (reluNode.inputNames[0] !== negInputNode.inputNames[0]) {
      continue;
    }

    const inputNode = reluNode.inputs[0];
    const outputNodes = addNode.children;

    // Construct a tensor for positive alpha (double negative).
    const alphaTensorName = negAlphaTensorNode.name + '_neg';

    const negNode: Node = {
      name: alphaTensorName,
      inputNames: [],
      inputs: [],
      attrParams: {},
      category: 'graph',
      children: [],
      op: 'Const',
      inputParams: {},
      rawAttrs: {}
    };

    // Add the constant to weightMap
    weightMap[alphaTensorName] = [neg(weightMap[negAlphaTensorNode.name][0])];
    graph.weights.push(negNode);

    // Construct the prelu node
    const preluNode: Node = {
      name: addNode.name + '_Prelu',
      inputNames: [inputNode.name, negNode.name],
      inputs: [inputNode, negNode],
      attrParams: {},
      category: 'custom',
      children: outputNodes,
      op: 'Prelu',
      inputParams: {
        'x': {inputIndexStart: 0, type: 'tensor'},
        'alpha': {inputIndexStart: 1, type: 'tensor'}
      }
    };

    negNode.children.push(preluNode);

    // Clean up the children and inputs of input/output nodes of the subgraph.
    const mulIndex = negAlphaTensorNode.children.indexOf(mulOp);
    if (mulIndex > -1) {
      negAlphaTensorNode.children.splice(mulIndex, 1);
    }

    const reluIndex = inputNode.children.indexOf(reluNode);
    if (reluIndex > -1) {
      inputNode.children.splice(reluIndex, 1);
    }

    const negIndex = inputNode.children.indexOf(negInputNode);
    if (negIndex > -1) {
      inputNode.children.splice(negIndex, 1);
    }
    inputNode.children.push(preluNode);

    outputNodes.forEach(node => {
      const addIndex = node.inputNames.indexOf(addNode.name);
      if (addIndex > -1) {
        node.inputNames[addIndex] = preluNode.name;
        node.inputs[addIndex] = preluNode;
      }
    });

    // The prelu node should be an output node.
    if (outputNodes.length === 0) {
      const addIndex = graph.outputs.indexOf(addNode);
      if (addIndex > -1) {
        graph.outputs.splice(addIndex, 1);
      }
      graph.outputs.push(preluNode);
    }
    // remove the nodes for keras generated prelu subgraph.
    delete graph.nodes[addNode.name];
    delete graph.nodes[mulOp.name];
    delete graph.nodes[reluNode.name];
    delete graph.nodes[reluNegInputNode.name];
    delete graph.nodes[negInputNode.name];

    // add the newly generated nodes.
    graph.nodes[preluNode.name] = preluNode;
    graph.nodes[negNode.name] = negNode;
  }
}
